//////////////////////////////////////////////////////////////////////////////////
// Who     : Lim, Dong-moon (dmlim@nate.com, 011-701-9381)
// When    : 2005/11/29
// History : 
//     - Applying for UNIX using OpenSSL
//////////////////////////////////////////////////////////////////////////////////

// NCrypto.h: interface for the CNCrypto class.
//
//////////////////////////////////////////////////////////////////////

#ifndef __NCRYPTO_H__
#define __NCRYPTO_H__

#ifdef WIN32

#if _MSC_VER > 1000
#pragma once
#endif // _MSC_VER > 1000

#include <windows.h>

#ifdef _WIN32_WCE
#else
#include <crtdbg.h>
#endif

#include <tchar.h>

#pragma warning( disable : 4786 )

#else

#include <ncrypt_types.h>

#endif // WIN32

#if defined(WIN32) && !defined(USE_OPENSSL)
#include <wincrypt.h>
#else
#include <openssl/evp.h>
#endif // defined(WIN32) && !defined(USE_OPENSSL)

#include <set>
#include <info_header.h>
#include <crypt_util.h>
#include <ncrypt_verify.h>

using namespace std;

//////////////////////////////////////////////////////////////////////////////////////

class CNCryptoContext;
class CNServerContext;
class CNCryptoSession;
class CNClientSession;
class CNServerSession;

namespace NCAPICOM
{
	class EncryptedData;
}

//////////////////////////////////////////////////////////////////////////////////////

class CNCryptoContext
{
public:
	CNCryptoContext();
	virtual ~CNCryptoContext();
	BOOL Init(BOOL isServer = FALSE);
	static void ClearAll();
	NBLOB* GetAlgorithmList();
	ALG_ID GetMostStrongAlgorithm(const NCRYPT_ALGORITHMS *pAlgs);

#if ((defined(__APPLE__) && defined(__MACH__)) || defined(WIN32)) && !defined(USE_OPENSSL)
#else
	LPCTSTR GetCertificateStore();
	void SetCertificateStore(LPCTSTR szCertStore);
#endif // WIN32

protected:

#if defined(WIN32) && !defined(USE_OPENSSL)
	HCRYPTPROV	GetProviderHandle();
	HCRYPTKEY	GetImportKey();
	HCRYPTPROV	m_hCryptProv;
	HCRYPTKEY	m_hImportKey;
#else
	LPTSTR m_szCertificateStore;
#endif // WIN32

	BOOL IsSupportedAlgorithm(ALG_ID algID);
	const NCRYPT_ALGORITHMS* GetAlgorithms();

	NCRYPT_ALGORITHMS m_algorithms;
	DWORD		m_cSupportAlgLen;
	ALG_ID*		m_pSupportAlgs;
	set<ALG_ID>	m_mapSupportAlgs;

	friend class CNCryptoSession;
	friend class CNServerSession;
	friend class CNClientSession;
	friend class CNDRMKeyContainer;
	friend class CNDRMSimpleCipher;
	friend class CNDRMEncrypter;
	friend class CNDRMDecrypter;
	friend class NCAPICOM::EncryptedData;
};

//////////////////////////////////////////////////////////////////////////////////////

#define MAX_CHAIN_CERT 10

class CNServerContext : public CNCryptoContext
{
public:
	CNServerContext();
	virtual ~CNServerContext();
	BOOL Init();
	BOOL ReadPvkAndCertFromFile(LPCTSTR pszPvkFile, LPCSTR pszPassword, LPCTSTR pszCertFile, LPCTSTR pszChainFile = NULL);
	const NBLOB* GetServerKeyInfo(DWORD blobType = NCRYPT_BLOB_BASE64);

#if defined(WIN32) && !defined(USE_OPENSSL)
	const HCRYPTKEY GetPrivateKeyHandle();
	HCRYPTKEY GetPrivateKeyHandleNew();
#else
	RSA*	GetRSAKey();
#endif // WIN32

protected:

	const NBLOB* GetServerCert();
	const NCRYPT_CERTIFICATES* GetServerChainCerts();

#if defined(WIN32) && !defined(USE_OPENSSL)
	HCRYPTKEY		m_hPrivateKey;
	TCHAR			m_szPvkFile[256];
	char			m_szPvkPassword[256];
#else
	RSA*			m_pRSAKey;
#endif // WIN32

	NBLOB			m_certBlob;
	NBLOB			m_keyInfoRaw;
	NBLOB			m_keyInfoB64;
	NBLOB			m_keyInfoB16;

	NCRYPT_CERTIFICATES m_chainCerts;
	NBLOB				m_chainsBlob[MAX_CHAIN_CERT];

	friend class CNServerSession;
};

//////////////////////////////////////////////////////////////////////////////////////

class CNCryptoSession
{
public:
	CNCryptoSession(CNCryptoContext* pContext);
	virtual ~CNCryptoSession();
	CNCryptoContext* GetContext();
	virtual BOOL Init(const BYTE*, DWORD, DWORD);
	NBLOB* EncryptSimple(const BYTE* pbPlain, DWORD cbPlain, DWORD blobType = NCRYPT_BLOB_RAW);
	NBLOB* DecryptSimple(const BYTE* pbEncrypted, DWORD cbEncrypted, DWORD blobType = NCRYPT_BLOB_RAW);
	NBLOB* EncryptMessage(const BYTE* pbPlain, DWORD cbPlain, DWORD blobType = NCRYPT_BLOB_RAW);
	NBLOB* DecryptMessage(const BYTE* pbEncrypted, DWORD cbEncrypted, DWORD blobType = NCRYPT_BLOB_RAW);
	BOOL   ImportPlainKey(ALG_ID algID, const BYTE* pbKey, DWORD cbKey, DWORD blobType = NCRYPT_BLOB_BASE16);
	NBLOB* ExportPlainKey(DWORD blobType = NCRYPT_BLOB_BASE16);
	NBLOB* GenerateCertMessage(LPCTSTR lpszMessage);
	BOOL   VerifyCertMessage(const BYTE* pCertMessage, DWORD cCertMessage, LPCTSTR lpszVerifyMessage);
	NBLOB* EncryptDRMKey(const BYTE* pbIn, DWORD cbIn, DWORD inBlobType = NCRYPT_BLOB_BASE64);
	ALG_ID GetAlgID();
	DWORD  GetPeerVersion();
	void   ClearPeerKeyInfo();

protected:
	NBLOB* GetCertHash(const NBLOB* pClientRandom, LPCTSTR lpszMessage);
	CNCryptoContext*	m_pContext;
	ALG_ID				m_cryptAlg;
	NBLOB				m_serverRandom;
	NCRYPT_KEYINFO*		m_pPeerKeyInfo;

#if defined(WIN32) && !defined(USE_OPENSSL)
	HCRYPTKEY			m_hSessionKey;
#else
	NBLOB				m_sessionKey;
#endif // WIN32
};

//////////////////////////////////////////////////////////////////////////////////////

class CNServerSession : public CNCryptoSession
{
public:
	CNServerSession(CNCryptoContext* pContext);
	virtual ~CNServerSession();
	virtual BOOL Init(const BYTE*, DWORD, DWORD blobType = NCRYPT_BLOB_BASE64, BOOL bAllowLowerVer = TRUE);
	const NBLOB* GetServerKeyInfo(DWORD blobType = NCRYPT_BLOB_BASE64, DWORD dwVersion = LATEST_NCRYPT_VER);
	BOOL ImportServerKeyInfo(const BYTE* pbKeyInfo, DWORD cbKeyInfo, DWORD blobType = NCRYPT_BLOB_BASE64);
	NBLOB* DecryptUsingPrivateKey(const NBLOB* pBlob, DWORD inBlobType = NCRYPT_BLOB_BASE64);

protected:
	NBLOB	m_keyInfoRaw;
	NBLOB	m_keyInfoB64;
	NBLOB	m_keyInfoB16;
};

//////////////////////////////////////////////////////////////////////////////////////

class CNClientSession : public CNCryptoSession
{
public:
	CNClientSession(CNCryptoContext* pContext);
	virtual ~CNClientSession();
	virtual BOOL Init(const BYTE* pbKey, DWORD cbKey, DWORD blobType = NCRYPT_BLOB_BASE64, BOOL bAllowLowerVer = TRUE);
	const NBLOB* ExportSessionKey(DWORD blobType = NCRYPT_BLOB_BASE64);
	void SetHostNameVerifier(HostNameVerifyCallback fn, LPVOID ctx = 0);
	NBLOB* EncryptUsingPublicKey(const NBLOB* pBlob, DWORD outBlobType = NCRYPT_BLOB_BASE64);
	BOOL VerifyServerCertificate(LPVOID pResult);
	BOOL ShowServerCertificate(LPVOID pData);
	void IgnoreNextSameCertError();

protected:

#if defined(WIN32) && !defined(USE_OPENSSL)
	HCRYPTKEY		m_hPublicKey;
#else
	X509*			m_pCert;
#endif // WIN32

	NBLOB			m_encKeyBlob;
	NBLOB			m_clientHash;

	NBLOB			m_keyInfoRaw;
	NBLOB			m_keyInfoB64;
	NBLOB			m_keyInfoB16;

	HostNameVerifyCallback	m_hostVerifier;
	LPVOID					m_hostVerifyCtx;

	BOOL			m_bIgnoreNextSameCertError;
	int				m_certsBackupLen;
	NBLOB			m_certsBackup[MAX_CHAIN_CERT+1];
};

//////////////////////////////////////////////////////////////////////////////////////

class CNDRMKeyContainer
{
public:
	CNDRMKeyContainer();
	~CNDRMKeyContainer();
	LPCTSTR InitRaw(CNCryptoContext* pContext, const BYTE* pbDrmKey, DWORD cbDrmSize);
	LPCTSTR Init(CNCryptoSession* pSession, const BYTE* pKey, DWORD cSize, DWORD blobType = NCRYPT_BLOB_BASE64);
	CNCryptoContext* GetContext();

protected:
	BOOL GetKeySet(const NBLOB* pOffsets, NBLOB* pKeySet);
	BOOL GetRandomKeySet(ALG_ID algId, NBLOB* pOffsets, NBLOB* pKeySet);

private:
	BOOL				DecryptTemp();
	void				ClearTemp();

	CNCryptoContext*	m_pContext;
	NBLOB				m_blobEncDRMKey;
	DWORD				m_cbDRMKey;
	NBLOB				m_blobTmpDRMKey;
	LPTSTR				m_szKeyHash;
	ALG_ID				m_algId;

#if defined(WIN32) && !defined(USE_OPENSSL)
	HCRYPTKEY			m_hCryptKey;
#else
	BYTE*				m_pCryptKey;
#endif

	friend class		CNDRMSimpleCipher;
	friend class		CNDRMEncrypter;
	friend class		CNDRMDecrypter;
};

//////////////////////////////////////////////////////////////////////////////////////

class CNDRMSimpleCipher
{
public:
	CNDRMSimpleCipher(CNDRMKeyContainer* pContainer);
	~CNDRMSimpleCipher();

	NBLOB* Encrypt(const char* algs, const BYTE* pData, DWORD cSize, DWORD blobType = NCRYPT_BLOB_BASE64);
	NBLOB* Decrypt(const BYTE* pData, DWORD cSize, DWORD blobType = NCRYPT_BLOB_BASE64);

private:
	CNDRMKeyContainer* m_pContainer;
};

//////////////////////////////////////////////////////////////////////////////////////

class CNDRMEncrypter
{
public:
	CNDRMEncrypter(CNDRMKeyContainer* pContainer);
	~CNDRMEncrypter();
	BOOL Init(const char* algs);
	DWORD GetBlockSize();
	DWORD AssumeHeaderSize();
	DWORD AssumeEncryptSize(DWORD dwSize);
	BOOL GetHeader(BYTE* pBuf, DWORD* inoutBufSize);
	BOOL EncryptUpdate(const BYTE *pInData, DWORD cInData, BYTE* pOutData, DWORD* pcOutData);
	BOOL EncryptFinal(const BYTE *pInData, DWORD cInData, BYTE* pOutData, DWORD* pcOutData);

private:
	CNDRMKeyContainer* m_pContainer;
	NDRM_CONTENT_HEADER m_drmHeader;

#if defined(WIN32) && !defined(USE_OPENSSL)
	HCRYPTKEY		m_hCryptKey;
#else
	EVP_CIPHER_CTX	m_cryptCtx;
#endif
};

//////////////////////////////////////////////////////////////////////////////////////

#define MAX_CONTENT_HEADER_SIZE 48

class CNDRMDecrypter
{
public:
	CNDRMDecrypter(CNDRMKeyContainer* pContainer);
	~CNDRMDecrypter();
	BOOL Init(BYTE* pbHeader, DWORD cbHeader, DWORD* pdwRead);
	DWORD GetBlockSize();
	DWORD AssumeDecryptSize(DWORD dwSize);
	BOOL DecryptUpdate(const BYTE* pInData, DWORD cInData, BYTE* pOutData, DWORD* pcOutData);
	BOOL DecryptFinal(const BYTE* pInData, DWORD cInData, BYTE* pOutData, DWORD* pcOutData);

private:
	CNDRMKeyContainer* m_pContainer;
	NDRM_CONTENT_HEADER* m_pDrmHeader;

#if defined(WIN32) && !defined(USE_OPENSSL)
	HCRYPTKEY		m_hCryptKey;
#else
	EVP_CIPHER_CTX	m_cryptCtx;
#endif

};

#endif // __NCRYPT_H__

extern const ALG_ID gPreferAlgOrder[];
